// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Authors: Alessandro Tasora, Radu Serban, Dario Fusai
// =============================================================================

#ifndef CHGNUPLOT_H
#define CHGNUPLOT_H

#include <sstream>
#include <iostream>
#include <iomanip>

#include "chrono/core/ChMatrix.h"
#include "chrono/assets/ChColor.h"
#include "chrono/functions/ChFunctionBase.h"
#include "chrono/functions/ChFunctionInterp.h"

#include "chrono_postprocess/ChApiPostProcess.h"

namespace chrono {

/// Namespace with classes for the postprocess unit.
namespace postprocess {

/// @addtogroup postprocess_module
/// @{

struct ChGnuPlotDataplot {
    ChMatrixDynamic<double> data;
    std::string command;
};

/// Class for plotting data with GNUplot.
/// This is a basic utility class which saves a temporary gpl file on disk and then calls the GNUplot utility from the
/// system shell, so it does not require to link GNUplot libraries. If the GNUplot is not installed, simply nothing
/// happens.
/// Notes:
/// - requires GNUplot version 4.6 or newer.
/// - the GNUplot executable must be available from a command window (make sure the executable inb is in the search
/// path).
/// - if no plot is displayed, open a cmd shell, type "gnuplot __tmp_gnuplot.gpl -persist" and see which error is
/// displayed.
class ChGnuPlot {
  public:
    ChGnuPlot(const std::string& filename = "__tmp_gnuplot.gpl") {
        m_commandfile += "# ------------------------------------------------------------------------------\n";
        m_commandfile += "# Autogenerated .gpl script that may be executed by gnuplot.exe\n";
        m_commandfile += "# Created by the ChGnuPlot class of the Chrono library\n";
        m_commandfile += "# ------------------------------------------------------------------------------\n\n";
        m_gpl_filename = filename;
        m_persist = true;
        m_map_plots[0] = {};
    }

    virtual ~ChGnuPlot() {
        FlushPlots(m_commandfile);
        ExecuteGnuplot(m_commandfile);
    }

    /// Enable/disable persistent plots (default: true).
    /// If false, the output window will close soon after the plot.
    /// For plotting in windows, it is better to use the default 'true' setting.
    /// When plotting to a file (EPS, PNG, PDF), this setting is usually set to 'false'.
    void SetPersist(bool state) { m_persist = state; }

    /// Add arbitrary GnuPlot commands in the gnuplot script.
    /// Basically you would just need calls to SetCommand() followed by an Output() at the end, however to make things
    /// easier, there are some shortcut functions such as SetRangeX(), SetGrid(), Plot(), etc. which populate the
    /// command file.
    void SetCommand(const std::string& command) { m_commandfile += command + "\n"; }

    /// This is equivalent to SetCommand().
    ChGnuPlot& operator<<(const std::string& command) {
        SetCommand(command);
        return *this;
    }

    /// Set default color sequence in plots.
    /// Available options:
    /// - 0: gnuplot default
    /// - 1: matplotlib
    /// - 2: matlab.
    void SetColorSequence(unsigned int num) {
        if (num == 0) {
            // Gnuplot default
            m_commandfile += "set linetype 1 lc rgb 'dark-violet'\n";
            m_commandfile += "set linetype 2 lc rgb 'sea-green'\n";
            m_commandfile += "set linetype 3 lc rgb 'cyan'\n";
            m_commandfile += "set linetype 4 lc rgb 'goldenrod'\n";
            m_commandfile += "set linetype 5 lc rgb 'dark-red'\n";
            m_commandfile += "set linetype 6 lc rgb 'blue'\n";
            m_commandfile += "set linetype 7 lc rgb 'dark-orange'\n";
            m_commandfile += "set linetype 8 lc rgb 'black'\n";
        } else if (num == 1) {
            // Matplotlib sequence
            m_commandfile += "set linetype 1 lc rgb '#1f77b4'\n";   // blue
            m_commandfile += "set linetype 2 lc rgb '#ff7f0e'\n";   // orange
            m_commandfile += "set linetype 3 lc rgb '#2ca02c'\n";   // green
            m_commandfile += "set linetype 4 lc rgb '#d62728'\n";   // red
            m_commandfile += "set linetype 5 lc rgb '#9467bd'\n";   // purple
            m_commandfile += "set linetype 6 lc rgb '#8c564b'\n";   // brown
            m_commandfile += "set linetype 7 lc rgb '#e377c2'\n";   // pink
            m_commandfile += "set linetype 8 lc rgb '#7f7f7f'\n";   // gray
            m_commandfile += "set linetype 9 lc rgb '#bcbd22'\n";   // olive
            m_commandfile += "set linetype 10 lc rgb '#17becf'\n";  // cyan
        } else if (num == 2) {
            // Matlab sequence
            m_commandfile += "set linetype 1 lc rgb '#0072bd'\n";  // blue
            m_commandfile += "set linetype 2 lc rgb '#d95319'\n";  // orange
            m_commandfile += "set linetype 3 lc rgb '#edb120'\n";  // yellow
            m_commandfile += "set linetype 4 lc rgb '#7e2f8e'\n";  // purple
            m_commandfile += "set linetype 5 lc rgb '#77ac30'\n";  // green
            m_commandfile += "set linetype 6 lc rgb '#4dbeee'\n";  // light blue
            m_commandfile += "set linetype 6 lc rgb '#a2142f'\n";  // dark red
        }

        FlushPlots();
    }

    /// Plot 2D (x,y) data from two dynamic Eigen vectors.
    void Plot(ChVectorDynamic<>& x,
              ChVectorDynamic<>& y,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        assert(x.size() == y.size());

        ChGnuPlotDataplot dataplot;
        dataplot.data.resize(x.size(), 2);
        dataplot.data.col(0) = x;
        dataplot.data.col(1) = y;

        dataplot.command +=
            "$x" + std::to_string(m_curr_idx) + std::to_string(m_map_plots.at(m_curr_idx).size()) + " using 1:2 ";

        dataplot.command += customsettings;
        dataplot.command += " title '" + title + "' ";

        m_map_plots[m_curr_idx].push_back(dataplot);
    }

    /// Plot 2D (x,y) data from two columns of a matrix.
    void Plot(ChMatrixConstRef data,
              int colX,
              int colY,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        ChVectorDynamic<> x = data.col(colX);
        ChVectorDynamic<> y = data.col(colY);
        Plot(x, y, title, customsettings);
    }

    /// Plot 2D (x,y) data from a ChFunctionInterp
    void Plot(ChFunctionInterp& fun_table,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        ChVectorDynamic<> x(fun_table.GetTable().size());
        ChVectorDynamic<> y(fun_table.GetTable().size());

        int i = 0;
        for (auto iter = fun_table.GetTable().begin(); iter != fun_table.GetTable().end(); ++iter) {
            x(i) = iter->first;
            y(i) = iter->second;
            ++i;
        }
        Plot(x, y, title, customsettings);
    }

    /// Plot 2D (x,y) data from two vectors.
    void Plot(const std::vector<double>& vals_x,
              const std::vector<double>& vals_y,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        ChVectorDynamic<> x = ChVectorDynamic<>::Map(vals_x.data(), vals_x.size());
        ChVectorDynamic<> y = ChVectorDynamic<>::Map(vals_y.data(), vals_y.size());
        Plot(x, y, title, customsettings);
    }

    /// Plot 2D (x,y) data from a generic ChFunction.
    /// Note that if the ChFunction is of type ChFunctionInterp there
    /// is a specific Plot() function
    void Plot(ChFunction& funct,
              double xmin,
              double xmax,
              double dx,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        int samples = static_cast<int>(std::floor((xmax - xmin) / dx));
        ChVectorDynamic<> x(samples);
        ChVectorDynamic<> y(samples);

        double x_crt = xmin;
        for (int i = 0; i < samples; ++i) {
            x(i) = x_crt;
            y(i) = funct.GetVal(x_crt);
            x_crt += dx;
        }
        Plot(x, y, title, customsettings);
    }

    /// Plot 2D (x,y) data from a generic ChFunction's 0th, 1st, 2nd, or 3rd derivative.
    void Plot(ChFunction& funct,
              int der_order,
              double xmin,
              double xmax,
              double dx,
              const std::string& title = "",
              const std::string& customsettings = " with lines ") {
        int samples = static_cast<int>(std::floor((xmax - xmin) / dx));
        ChVectorDynamic<> x(samples);
        ChVectorDynamic<> y(samples);

        double x_crt = xmin;
        for (int i = 0; i < samples; ++i) {
            x(i) = x_crt;
            if (der_order == 0)
                y(i) = funct.GetVal(x_crt);
            else if (der_order == 1)
                y(i) = funct.GetDer(x_crt);
            else if (der_order == 2)
                y(i) = funct.GetDer2(x_crt);
            else if (der_order == 3)
                y(i) = funct.GetDer3(x_crt);
            else
                y(i) = 0;
            x_crt += dx;
        }
        Plot(x, y, title, customsettings);
    }

    /// Plot 2D (x,y) data from an external file.
    void Plot(const std::string& datfile,
              int colX,
              int colY,
              const std::string& title,
              const std::string& customsettings = " with lines ") {
        ChGnuPlotDataplot dataplot;

        dataplot.command += " \"";
        dataplot.command += datfile;
        dataplot.command += "\" using ";
        dataplot.command += std::to_string(colX);
        dataplot.command += ":";
        dataplot.command += std::to_string(colY);
        dataplot.command += " ";
        dataplot.command += customsettings;
        dataplot.command += " title \"";
        dataplot.command += title;
        dataplot.command += "\" ";

        m_map_plots[m_curr_idx].push_back(dataplot);
    }

    /// Start subplot in current figure.
    /// Overall subplot layout is provided as
    /// - total number of rows
    /// - total number of columns
    /// - index of current subplot.
    /// Notes:
    /// - indices must be given in ascending order (starting from 0)
    /// - indices must cover all specified subplots
    /// - requires manual call to EndSubplot() when all subplots have been completed.
    void StartSubplot(int rows, int cols, int idx) {
        if (idx < 0 || idx > rows * cols - 1 || idx < m_curr_idx)
            return;

        FlushPlots();
        m_curr_idx = idx;

        if (m_curr_idx == 0) {
            m_subplots = true;
            m_end_idx = rows * cols;
            m_commandfile += "set multiplot layout " + std::to_string(rows) + "," + std::to_string(cols) + "\n";
        } else {
            m_map_plots[m_curr_idx] = {};
        }
    }

    /// End current subplot layout.
    /// Call this after all plots set by StartSubplot() have been completed.
    void EndSubplot() {
        if (!m_subplots)
            return;

        FlushPlots();
        m_commandfile += "unset multiplot\n\n";
        m_subplots = false;
        m_curr_idx = 0;
        m_end_idx = 0;
    }

    /// Add a plot title.
    void SetTitle(const std::string& label) { m_commandfile += "set title '" + label + "'\n"; }

    /// Add a label to the X axis.
    void SetLabelX(const std::string& label) { m_commandfile += "set xlabel '" + label + "'\n"; }

    /// Add a label to the Y axis.
    void SetLabelY(const std::string& label) { m_commandfile += "set ylabel '" + label + "'\n"; }

    /// Add a label to the Y2 axis.
    void SetLabelY2(const std::string& label) { m_commandfile += "set y2label '" + label + "'\n"; }

    /// Add a label to the Z axis.
    void SetLabelZ(const std::string& label) { m_commandfile += "set zlabel '" + label + "'\n"; }

    /// Show legend and set optional parameters.
    /// Example: SetLegend("bottom right box opaque")
    void SetLegend(const std::string& customsettings) { m_commandfile += " set key " + customsettings + "\n"; }

    /// Hide plot legend.
    void HideLegend() { m_commandfile += " unset key\n"; }

    /// Enable and specify parameters for a plot grid.
    void SetGrid(bool dashed = true, double linewidth = 1.0, const ChColor& color = ChColor(0, 0, 0)) {
        m_commandfile += "set grid ";
        m_commandfile += dashed ? "lt 0 " : "lt 1 ";
        m_commandfile += "lw " + std::to_string(linewidth);
        m_commandfile += " lc " + col_to_hex(color);
        m_commandfile += "\n";
    }

    /// Set the X data range.
    void SetRangeX(double xmin, double xmax, bool automin = false, bool automax = false) {
        m_commandfile += "set xrange [";
        m_commandfile += automin ? "*" : std::to_string(xmin);
        m_commandfile += ":";
        m_commandfile += automax ? "*" : std::to_string(xmax);
        m_commandfile += "] \n";
    }

    /// Set the Y data range.
    void SetRangeY(double ymin, double ymax, bool automin = false, bool automax = false) {
        m_commandfile += "set yrange [";
        m_commandfile += automin ? "*" : std::to_string(ymin);
        m_commandfile += ":";
        m_commandfile += automax ? "*" : std::to_string(ymax);
        m_commandfile += "] \n";
    }

    /// Set the Y2 data range.
    void SetRangeY2(double ymin2, double ymax2, bool automin = false, bool automax = false) {
        m_commandfile += "set y2range [";
        m_commandfile += automin ? "*" : std::to_string(ymin2);
        m_commandfile += ":";
        m_commandfile += automax ? "*" : std::to_string(ymax2);
        m_commandfile += "] \n";
    }

    /// Set the Z data range.
    void SetRangeZ(double zmin, double zmax, bool automin = false, bool automax = false) {
        m_commandfile += "set zrange [";
        m_commandfile += automin ? "*" : std::to_string(zmin);
        m_commandfile += ":";
        m_commandfile += automax ? "*" : std::to_string(zmax);
        m_commandfile += "] \n";
    }

    /// Set aspect ratio.
    void SetAspectRatio(double val) { m_commandfile += " set size ratio " + std::to_string(val) + "\n"; }

    /// Set axes to equal size (i.e., square box) or restore default.
    void SetAxesEqual(bool axequal) { m_commandfile += axequal ? " set size square \n" : " set size nosquare \n"; }

    /// Set output window title.
    void SetOutputWindowTitle(const std::string& label) { m_commandfile += "set term wxt title '" + label + "'\n"; }

    /// Set canvas size.
    void SetCanvasSize(int width, int height) {
        m_commandfile += "set term wxt size ";
        m_commandfile += std::to_string(width);
        m_commandfile += ",";
        m_commandfile += std::to_string(height);
        m_commandfile += "\n";
    }

    /// Replot last plots (since the last time any of the functions Output*** was used).
    void Replot() { m_commandfile += "replot \n"; }

    /// Set plot in an interactive window.
    /// For multiple windows, call this with increasing windownum, interleaving with Plot() statements etc.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputWindow(int windownum = 0, bool reset_text = false) {
        FlushPlots(m_commandfile);
        m_commandfile += "# Figure " + std::to_string(windownum) + "\n";
        m_commandfile += "set terminal wxt " + std::to_string(windownum) + "\n";

        if (reset_text) {
            SetLabelX("");
            SetLabelY("");
            SetLabelY2("");
            SetLabelZ("");
            SetTitle("");
        }
    }

    /// Save plot in a PNG file.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputPNG(const std::string& filename, int sizex = 400, int sizey = 300) {
        FlushPlots(m_commandfile);
        m_commandfile += "set terminal png size ";
        m_commandfile += std::to_string(sizex);
        m_commandfile += ",";
        m_commandfile += std::to_string(sizey);
        m_commandfile += "\n";
        m_commandfile += "set output '";
        m_commandfile += filename;
        m_commandfile += "'\n";
    }

    /// Save plot in a PDF file.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputPDF(const std::string& filename) {
        FlushPlots(m_commandfile);
        m_commandfile += "set terminal pdf\n";
        m_commandfile += "set output '";
        m_commandfile += filename;
        m_commandfile += "'\n";
    }

    /// Save plot in a EPS file.
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputEPS(const std::string& filename, double inchsizex = 4, double inchsizey = 3, bool color = true) {
        FlushPlots(m_commandfile);
        m_commandfile += "set terminal postscript eps ";
        if (color)
            m_commandfile += " color ";
        m_commandfile += " size ";
        m_commandfile += std::to_string(inchsizex);
        m_commandfile += ",";
        m_commandfile += std::to_string(inchsizey);
        m_commandfile += "\n";
        m_commandfile += "set output '";
        m_commandfile += filename;
        m_commandfile += "'\n";
    }

    /// Save plot in a custom terminal.
    /// For instance try terminalsetting ="set terminal svg size 350,262 fname 'Verdana' fsize 10"
    /// Call this before Plot() statements. Otherwise call Replot() just after.
    void OutputCustomTerminal(const std::string& filename, const std::string& terminalsetting) {
        FlushPlots(m_commandfile);
        m_commandfile += terminalsetting;
        m_commandfile += "set output '";
        m_commandfile += filename;
        m_commandfile += "'\n";
    }

    /// Close current plot command (e.g., for use in multiplot).
    void FlushPlots() { FlushPlots(m_commandfile); }

  protected:
    void FlushPlots(std::string& script) {
        for (int idx = 0; idx < m_map_plots.size(); ++idx) {
            // Embed plot data in the .gpl file
            for (int i = 0; i < m_map_plots.at(idx).size(); ++i) {
                if ((m_map_plots.at(idx)[i].data.cols() > 0) && (m_map_plots.at(idx)[i].data.rows() > 0)) {
                    script += "$x" + std::to_string(idx) + std::to_string(i) + " << EOD\n";

                    for (int ir = 0; ir < m_map_plots.at(idx)[i].data.rows(); ++ir) {
                        for (int ic = 0; ic < m_map_plots.at(idx)[i].data.cols(); ++ic) {
                            script += std::to_string(m_map_plots.at(idx)[i].data(ir, ic));
                            script += " ";
                        }
                        script += "\n";
                    }

                    script += "EOD\n";
                }
            }

            if (m_map_plots.at(idx).size() > 0)
                script += "plot ";

            for (int i = 0; i < m_map_plots.at(idx).size(); ++i) {
                if (i > 0)
                    script += " ,\\\n     ";
                script += m_map_plots.at(idx)[i].command;
            }
            // script += "\n";

            m_map_plots.at(idx) = {};
            // m_map_plots.at(idx).resize(0);
        }
        script += "\n";
    }

    void ExecuteGnuplot(std::string& script) {
        // Create a tmp .gpl file
        {
            std::ofstream gnuplot_command(m_gpl_filename);
            gnuplot_command << script;
        }

        std::string syscmd;

        // Launch the GNUplot from shell
#ifdef _WIN32
        // "start /b gnuplot __tmp_gnuplot.gpl -persist"
        // where /b avoids showing the black cmd window
        syscmd += "start /b gnuplot \"";
        syscmd += m_gpl_filename;
        syscmd += "\"";
        if (m_persist)
            syscmd += " -persist";
        system(syscmd.c_str());
#else
        // Unix like systems:
        // "gnuplot __tmp_gnuplot.gpl -persist &"
        syscmd += "gnuplot \"";
        syscmd += m_gpl_filename;
        syscmd += "\"";
        if (m_persist)
            syscmd += " -persist";
        syscmd += " &";  // to launch and forget
        system(syscmd.c_str());
#endif
    }

    std::string col_to_hex(ChColor color) {
        std::string ret;
        ret += "rgb '#";
        std::stringstream streamR;
        streamR << std::setfill('0') << std::setw(2) << std::hex << static_cast<int>(color.R * 255);
        std::string sr(streamR.str());
        ret += sr;
        std::stringstream streamG;
        streamG << std::setfill('0') << std::setw(2) << std::hex << static_cast<int>(color.G * 255);
        std::string sg(streamG.str());
        ret += sg;
        std::stringstream streamB;
        streamB << std::setfill('0') << std::setw(2) << std::hex << static_cast<int>(color.B * 255);
        std::string sb(streamB.str());
        ret += sb;
        ret += "'";
        return ret;
    }

    std::string m_gpl_filename;  ///< gnuplot file name
    std::string m_commandfile;   ///< gnuplot script
    bool m_persist = true;       ///< prevent window closure after plotting

    bool m_subplots = false;  ///< activate subplot mode
    int m_curr_idx = 0;       ///< current subplot index
    int m_end_idx = 0;        ///< index of last subplot
    std::unordered_map<int, std::vector<ChGnuPlotDataplot>>
        m_map_plots;  ///< index of subplot in figure, multiplots inside subplot
};

/// @} postprocess_module

}  // end namespace postprocess
}  // end namespace chrono

#endif